# CONTEXT {'gpu_vendor': 'AMD', 'guest_os': 'UBUNTU'}
# FBGEMM requires pytorch 1.10
ARG BASE_DOCKER=compute-artifactory.amd.com:5000/rocm-plus-docker/framework/compute-rocm-dkms-no-npi-hipclang:12471_ubuntu20.04_py3.8_pytorch_release-2.0_e0e83ab

FROM $BASE_DOCKER
USER root
ENV WORKSPACE_DIR=/workspace
RUN mkdir -p $WORKSPACE_DIR
WORKDIR $WORKSPACE_DIR

# ROCm gpg key
#RUN wget -q -O - http://repo.radeon.com/rocm/rocm.gpg.key | sudo apt-key add -
#RUN apt update && apt install -y \
#    unzip \
#    jq

# add sshpass, sshfs for downloading from mlse-nas
#RUN apt-get install -y sshpass sshfs
#RUN apt-get install -y netcat

#RUN apt install net-tools

# add locale en_US.UTF-8
#RUN apt-get install -y locales
#RUN locale-gen en_US.UTF-8
#RUN pip3 install --upgrade pip

# Remove triton and install new
#RUN apt remove triton -y
RUN pip uninstall triton -y
RUN git clone https://github.com/ROCmSoftwarePlatform/triton && cd triton && git checkout 6b64ef959d41b7b229dafe7c1112b1f3753a2aa7 && cd python && python setup.py develop && cd ../..

# Install DeepSpeed (https://github.com/ROCmSoftwarePlatform/DeepSpeed/blob/master/docker/Dockerfile.rocm)
RUN pip3 install deepspeed